package org.stragma.apply.modules

import org.kiama.rewriting._
import org.kiama.rewriting.Rewriter.{all,attempt,debug,fail,id,rulefs,Strategy,Term,where}

import org.stragma.apply.signatures.Apply._
import org.stragma.apply.signatures.Options._
import org.stragma.apply.signatures.Conversions._
import org.stragma.apply.modules._
import org.stragma.apply.modules.Addassertions.{add_assertions, remove_duplicate_assertions}
import org.stragma.apply.modules.Beautify._
import org.stragma.apply.modules.Normalise.{normalise,normaliseTest}
import org.stragma.apply.modules.Options._

import org.stragma.apply.util._
import org.stragma.apply.util.LogicalOperators._

import org.stragma.util.Configuration._

/**
 *	This class is an implementation os the Simplification process used in the
 *	Apply compilation pipeline
 */
class Simplify {
	
	// Initialise global rule to fail until set
	var propagateConstant : DynamicRules = new DynamicRules
	val assertions = new Assertions
	
	/**
	 *	This strategy was ported from Apply, though it does look strange it is used
	 *	in the strategy propagate_constant_assertion below
	 */
	def debug_simplify = id
	
	/*
	 * This helper method execute simplify "n" number of times
	 */
	def simplifyHelper ( term : Term, n : scala.Int ) : Term = {
		var result : Term = scala.None
		for(i <- 0 until n) {
			result = (simplify(term)).get
			
			// Added the println to mimic the way strategoxt.jar works
			println(result)
		}
		
		result
	}
	
	/*
	 * This helper method executes normaliseTest n times and returns
	 * the very last result to the caller
	 */
	def simplifyHelper ( file : String, n : scala.Int ) : Term = {
		val iTerm = loadTermFromFile (file)
		simplifyHelper (iTerm, n)
	}
	
	/**
	 * Port of primary strategy from Simplify
	 */
	def simplify = 
		(add_assertions <+ where(debug("add-assertions failed\n"))) <*
		(propagate_constant <+ where(debug("propagate-constant failed\n"))) <*
		(beautify1 <+ where(debug("beautify failed\n"))) <*
		(remove_duplicate_assertions <+ where(debug("remove-duplicate-assertions failed\n")))
	
	/**
	 * First apply the dynamic PropagateConstant rules to the current tree, 
   	 * but if that fails see whether it is an assignment that will generate a dynamic rule,
   	 * and if that fails propagate constants in all sub trees and then simplify the current tree.
   	 * Note that a propagated constant is always simple, so does not need to be further simplified.
   	 * If constant propagation does not generate a value for the node, propagate constants
   	 * in all sub nodes, then try simplifying the node.
   	 * One interesting issue: assertions mostly match the comparison expression,
   	 * so they can be tried directly on the normalised comparison expression
  	 * as propagate-constant climbs back up the tree.
  	 * However, equality assertions do not have this property - they match the
  	 * contents of a comparison (or other expression) that has been normalised.
  	 * So, if normalisation changes the tree then we must try equality assertions
  	 * recursively in the tree and if anything happens there then we should 
  	 */
 	private def propagate_constant_part1 : Strategy = {
		propagateConstant.getStrategies <+
		propagate_constant_assign <+
		propagate_constant_assert <+
		propagate_constant_cassign <+
		propagate_constant_if <+
		propagate_constant_ifknown <+
		propagate_constant_while <+
		propagate_constant_for <+
		propagate_constant_cfor <+
		propagate_constant_procedure <+
		propagate_constant_declarations
	}
	
	def propagate_constant : Strategy = {
		all(propagate_constant_part1) <+
		attempt(normalise) <*
		assertions.try_assertion
	}
	
	/**
	 * This version of propagate constant was used to debug the strategy, this strategy is not used in the performance tests
	 */
	def propagate_constant_debug : Strategy = {
		( propagateConstant.getStrategies <+ where(debug("PropagateConstant failed ->")) <* fail ) <+
		( propagate_constant_assign <+ where(debug("propagate_constant_assign failed ->")) <* fail ) <+
		( propagate_constant_assert <+ where(debug("propagate_constant_assert failed ->")) <* fail ) <+
		( propagate_constant_cassign <+ where(debug("propagate_constant_cassign failed ->")) <* fail ) <+
		( propagate_constant_if <+ where(debug("propagate_constant_if failed ->")) <* fail ) <+
		( propagate_constant_ifknown <+ where(debug("propagate_constant_ifknown failed ->")) <* fail ) <+
		( propagate_constant_while <+ where(debug("propagate_constant_while failed ->")) <* fail ) <+
		( propagate_constant_for <+ where(debug("propagate_constant_for failed ->")) <* fail ) <+
		( propagate_constant_cfor ) <+
		( propagate_constant_procedure <+ where(debug("propagate_constant_procedure failed ->")) <* fail ) <+
		( propagate_constant_declarations <+ where(debug("propagate_constant_declarations failed ->")) <* fail ) <+
		( all(propagate_constant_debug) <+ where(debug("all(propagate_constant) failed ->")) <* fail ) <* 
		( attempt(normaliseTest) <+ where(debug("attempt normalise failed ->")) <* fail ) <* 
		( assertions.try_assertion <+ where(debug("try assertion failed ->")) )
	}
	
	/**
	 * Port of the strategy propogate-constant-assign from Simplify.str
	 */
	def propagate_constant_assign = rulefs {
		case Assign ( TVarToTScalarVar(Var(x)), y ) =>
			
			// We have to use an attempt here because a Stratego command failure and Kiama one is not the same thing
			val e = attempt( propagate_constant(y) )
			
			def initPropagateConstant = {
				if (get_config(Propagate_constant_assignments()).isDefined && isValue(e).isDefined) {
					propagateConstant + ( Var(x), rulefs { case x : Var => e } )
				}
				else
					propagateConstant - Var(x)
			}
			
			where(assertions.undefine_assertions(x)) <* initPropagateConstant
	}	

	/**
	 * constant propagation inside assert statements
  	 * works by breaking down the assert statement into
     * smaller parts where possible.
	 */
	def propagate_constant_assert =
		rulefs {
			case Assert( x ) =>
				where ( Assert(attempt(propagate_constant_assertion(x))) ) <+ id
		}
	
	 /**
	  * An assertion that is an and operation is treated in parts
      * An 'ordinary' assertion: propagate constants inside it 
      * (includes normalising it), then use it to generate rules
	  */
	def propagate_constant_assertion : Strategy =
		rulefs {
			case And(x,y) =>
				And(attempt(propagate_constant_assertion(x)), attempt(propagate_constant_assertion(y)))
		} <+ 
		debug_simplify <*
		rulefs {
			case a =>
				propagate_constant <* assertions.define_assertion(a)
		}

	/**
	 * Port of the strategy propogate-constant-cassign from Simplify.str
	 */
	def propagate_constant_cassign = 
		rulefs {
			case CAssign ( Var(x), y ) =>
				val e = attempt ( propagate_constant(y) )
			
				def initPropagateConstant = {
					// We use isDefined because the existence of a value is Scala is not implicit like
					// in Stratego, in addition our strategies return a Scala Option[T] type which
					// supports the isDefined method
					if (get_config( Propagate_constant_assignments() ).isDefined && isValue( e ).isDefined)
						propagateConstant + ( Var(x), rulefs { case Var(x) => e } )
					else
						propagateConstant - Var(x)
				}
			
			where(assertions.undefine_assertions(x)) <* initPropagateConstant
	}	

	/**
	 * Port of the strategy propagate-constant-if from Simplify, note that the multiple
	 * definitions from Stratego have been merged into one because Scala supports the
	 * implementation of multiple matches within single partial functions.
	 */
	def propagate_constant_if = 
		rulefs {
			case If (cond, stmts) =>
				If ( attempt(propagate_constant(cond)) , stmts)
			case IfElse(cond, trueBit, falseBit) =>
				IfElse(attempt(propagate_constant(cond)), trueBit, falseBit)
			case IfElseExp(cond, trueBit, falseBit) =>
				IfElseExp(attempt(propagate_constant(cond)), trueBit, falseBit)
		} <* 
		evalIf <* 
		propagate_constant <+ 
		rulefs {
			case If(cond, statements) =>
				If(cond, attempt(propagate_constant(statements)))
			case IfElse(cond, trueBit, falseBit) =>
				IfElse(cond, attempt(propagate_constant(trueBit)), attempt(propagate_constant(falseBit)))
			case IfElseExp(cond, trueBit, falseBit) =>
				IfElseExp(cond, attempt(propagate_constant(trueBit)), attempt(propagate_constant(falseBit)))
		}

	/**
	 * Note that solve-ifknown must always succeed by definition
	 */
	def propagate_constant_ifknown = 
		rulefs {
			case IfKnownExp ( cond, trueBit, falseBit ) =>
				IfKnownExp ( attempt(propagate_constant(cond)), trueBit, falseBit )
		} <* 
		( 
			solveIfKnown <* 
			propagate_constant 
		)
	
	/**
	 * Propagation of constant in a while loop.
   	 * The first line simply matches the while loop -- there are no parts of the loop
  	 * that are unconditionally evaluated once only.
   	 * A second line tries propagating constants into the loop condition, and if this
   	 * allows the while loop to be eliminated, then we are done.
   	 * Otherwise, the last line starts with the original while loop and propagates constants
   	 * in the condition and the body until convergence is reached.
	 */
	def propagate_constant_while = 
		rulefs {
			case While ( cond, body ) =>
				While ( attempt( propagate_constant(cond) ), body )
		} <* 
		evalWhile <* 
		propagate_constant <+ 
		rulefs {
			case While ( cond, body ) =>
				While ( attempt( propagate_constant(cond) ), attempt( propagate_constant(body) ) )
		}
	
	/**
	 * A for loop has the range evaluated before the loop is entered.
  	 * If this allows the loop to be eliminated, then we are done.
  	 * Otherwise, we must propagate constants in the body.
  	 * The loop variable itself is not constant.
  	 */
	def propagate_constant_for = 
		rulefs {
			case For ( Var(x), from, to , statements ) =>
				propagateConstant - Var(x)
				assertions undefine_assertions ( Var(x) )
				id
		} <*
		rulefs {
			case For ( varname, from, to, statements ) =>
				For ( varname, attempt( propagate_constant(from) ), attempt( propagate_constant(to) ), statements )
		} <* rulefs {
			case For ( varname, from, to, statements ) =>
				For ( varname, from, to, attempt( propagate_constant(statements) ) )
		} <* ( evalFor <* propagate_constant )
	
	/**
	 * A C style for loop has the initialisation performed before the loop is entered.
  	 * Similar to the while loop, we first try to eliminate the loop.
  	 * If that fails, we must propagate constants in the condition, the loop body
  	 * and the increments.  The increments are done last, so they must have the rule applied last.
  	 * That is why there are two applications in the final iteration.
	 */
	def propagate_constant_cfor =
		rulefs {
			case CFor ( init, cond, incs, stmts ) =>
				CFor ( attempt(propagate_constant(init)), cond, incs, stmts )
		} (
			rulefs {
				case CFor ( init, cond, incs, stmts ) =>
				CFor ( init, attempt(propagate_constant(cond)), incs, stmts )
			} <*
			evalCFor <* 
			propagate_constant <+ 
			(
				rulefs {
					case CFor ( init, cond, incs, stmts ) =>
						CFor ( init, attempt(propagate_constant(cond)), incs, attempt(propagate_constant(stmts)) )
				} <*
				rulefs {
					case CFor ( init, cond, incs, stmts ) =>
						CFor ( init, cond, attempt(propagate_constant(incs)), stmts )
				}
			)
		)
	
	/**
	 * Limit scope of constant propagation to a single procedure at a time
   	 * We have no nested scopes (at the moment) so we don't need labelled scopes.
	 */
	def propagate_constant_procedure = 
		rulefs {
			case Procedure ( name, args, decs, stmts, ename ) =>
				Procedure ( name, args, attempt(propagate_constant(decs)), attempt(propagate_constant(stmts)), ename )
		}
	
	/**
	 * These are the If rules in Simplify.str
	 */
	val evalIf = 
		rulefs {
			// EvalIf : If(False(), x)             		-> Statements([])
			// EvalIf : IfElse(False(), x, y)      		-> y
			// EvalIf : IfElseExp(False(), x, y)   		-> y
			// EvalIf : If(True(), x)              		-> x
			// EvalIf : IfElse(True(), x, y)       		-> x
			// EvalIf : IfElseExp(True(), x, y)    		-> x
			case If ( False(), x ) 						=> Statements(Nil)
			case If ( True(), x )						=> x
			case IfElse ( False(), x, y )				=> y
			case IfElse ( True(), x, y )				=> x
			case IfElseExp ( False(), x, y )   			=> y
			case IfElseExp ( True(), x, y )    			=> x
		}
	
	/*
	 * Stratego rule from Simplify.str
	 */
	val evalWhile = 
		rulefs {
			// EvalWhile : While(False(), x)       			-> Statements([])
			case While ( False(),x )						=> Statements(Nil)
		}
	
	/*
	 * Stratego rules from Simplify.str
	 */
	val evalIfKnownTrue = 
		rulefs {
			// EvalIfKnownTrue : IfKnownStmt(True, x, y) 	-> x
			// EvalIfKnownTrue : IfKnownExp(True, x, y)  	-> x
			case IfKnownStmt ( True(),x,y )					=> x
			case IfKnownExp ( True(),x,y )					=> x
		}
	
	/*
	 * Stratego rules from Simplify.str
	 */
	val evalIfKnownOther = 
		rulefs {
			// EvalIfKnownOther : IfKnownStmt(c, x, y)   	-> y
			// EvalIfKnownOther : IfKnownExp(c, x, y)    	-> y
			case IfKnownStmt ( c,x,y )						=> y
			case IfKnownExp ( c,x,y )						=> y
		}
	
	/**
	 * Stratego rules from Simplify.str
	 */
	val evalFor = 
		rulefs {
			// EvalFor : For(x, a, a, z)           					-> Statements([Assign(x,a), z])
			// EvalFor : For(x, Int(i), Int(j), z) 					-> Statements([Assign(x,Int(i))]) where <ltS>(j,i)
			// Needs to be defined first as Int is an Exp
			case For (x, i:Int, j:Int, z ) if (ltS(j,i).isDefined) 	=> Statements ( Assign(TVarToTScalarVar(x),i) :: Nil )
			case For (x, a, b, Statements(z) ) if (a==b) 			=> Statements ( List(Assign(TVarToTScalarVar(x),a)) ::: z )
		}
	
	/**
	 * Stratego rules from Simplify.str
	 */
	val evalCFor = 
		rulefs {
			// EvalCFor : CFor(init, False(), inc, s) 	-> Statements([Statement(init)])  // where <debug>["EvalCFor: ", init]
			case CFor ( init, False(), inc, s ) 		=> Statements ( Statement(init) :: Nil )
		}
	
	/**
	 * Port of the solve-ifknown strategy
	 */
	val solveIfKnown = 
		evalIfKnownTrue <+ evalIfKnownOther
	
	/**
	 * Port of is-value strategy
	 * is-value = ?Int(i)
	 */
	val isValue = 
		rulefs { case Int(i) => id }
	
	/**
	 * propagate-constant-declarations = ?Declarations(x)
	 */
	val propagate_constant_declarations = 
		rulefs { case Declarations(x) => id }
}